print_df(head(df))

df_apim <- df %>%
  # Reshaping to Actor-Partner format (4-field)
  reshape_dyadic_data(
    person_id = 'userID',
    dyad_id = 'coupleID',
    vars_to_reshape = 'communication'
  )  %>%
  mutate(
    # Optional: grand-mean centering
    communication_actor_gmc = communication_actor - mean(communication_actor),
    communication_partner_gmc = communication_partner - mean(communication_partner),
    
    # Create Dummy-Variables for male and female
    is_female = ifelse(gender == 1, 1, 0),
    is_male = ifelse(gender == 2, 1, 0),
  ) 

print_df(head(df_apim))

# Option A: Random-intercept model for non-independence

formula <- bf(
  satisfaction ~ 
    
    # Remove global intercept, introduce male and female intercepts.
    0 + is_male + is_female +
    
    # Actor effect for male and female
    communication_actor_gmc:is_male + 
    communication_actor_gmc:is_female + 
    
    # Partner effect for male and female
    communication_partner_gmc:is_male + 
    communication_partner_gmc:is_female +
    
    # Option 0: Random intercept for male and female (correlated)
    # NOT identifiable in the cross-sectional case:
    # (0 + is_male + is_female | coupleID)
    
    # Option A: Shared dyad effect - interpretable dyad-level random intercept. 
    (1 | coupleID), 
  
  # Two distinct residual variances for males vs females:
  sigma = ~ 0 + is_male + is_female
)

priors <- c(
  # prior(normal(2, 3), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "b", dpar = "sigma"),
  prior(student_t(3, 0, 2.5), class = "sd")
)

model_dist_apim_a <- brm(
  formula = formula, 
  data = df_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_dist_apim_a') # Cache the model
)

# Option B: Residual CS for non-independence (no random intercept)

formula <- bf(
  satisfaction ~ 
    
    # Remove global intercept, introduce male and female intercepts.
    0 + is_male + is_female +
    
    # Actor effect for male and female
    communication_actor_gmc:is_male + 
    communication_actor_gmc:is_female + 
    
    # Partner effect for male and female
    communication_partner_gmc:is_male + 
    communication_partner_gmc:is_female +
  
    # Option B: Model non-independence by using compound symmetry of the residuals. 
    # Aligns with the SEM conceptualization / literature
    cosy(gr = coupleID),
  
  # Two distinct residual variances for males vs females:
  sigma = ~ 0 + is_male + is_female
)

priors <- c(
  # prior(normal(2, 3), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "b", dpar = "sigma"),
  prior(beta(1, 3), class = "cosy")
)

model_dist_apim_b <- brm(
  formula = formula, 
  data = df_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_dist_apim_b') # Cache the model
)

rstan::check_hmc_diagnostics(model_dist_apim_b$fit)

loo::pareto_k_table(loo(model_dist_apim_b))

plot(model_dist_apim_b, ask = FALSE)

pp_check(model_dist_apim_b, 'dens_overlay_grouped', group = 'is_male')

pp_check(model_dist_apim_b, 'ecdf_overlay_grouped', group = 'is_male')

pp_check(model_dist_apim_b, type = "loo_pit_overlay")

# Custom function to make DHARMa work with brms (see file 'Functions')
DHARMa.check_brms(model_dist_apim_b)

df_long_apim <- bmlm::isolate(
  df_long_apim, 
  by = 'userID', # NOT coupleID
  value = c('provided_support_actor', 'provided_support_partner'),
  which = 'both' 
) %>%
# renaming to avoid confusion with between- and within COUPLE centred DIM variables
  rename(
    provided_support_actor_cwp = provided_support_actor_cw,
    provided_support_partner_cwp = provided_support_partner_cw,
    provided_support_actor_cbp = provided_support_actor_cb,
    provided_support_partner_cbp = provided_support_partner_cb,
  ) %>%
  # Centering time
  mutate(
    diaryday_c = scale(diaryday, center = TRUE, scale = FALSE)
  )

library(nlme)
# nlme model for reference
model1 <- lme(
  fixed = 
    # Intercepts
    closeness ~ 0 + is_male + is_female +
    
    # Time Slopes
    is_male:diaryday_c +
    is_female:diaryday_c +
    
    # Between-Person APIM
    is_male:provided_support_actor_cbp + 
    is_male:provided_support_partner_cbp +
    
    is_female:provided_support_actor_cbp + 
    is_female:provided_support_partner_cbp +
  
    # Within-Person APIM
    is_male:provided_support_actor_cwp + 
    is_male:provided_support_partner_cwp +
    
    is_female:provided_support_actor_cwp + 
    is_female:provided_support_partner_cwp,
  
  random = 
    ~ 0 + is_male + is_female 
       
    + is_male:diaryday_c + is_female:diaryday_c
   
  # in nlme we need to remove many random slopes to achieve convergence
  # alternatively, we could try to remove correlations between the slopes. 
    #+ is_male:provided_support_actor_cwp 
    #+ is_male:provided_support_partner_cwp
    #+ is_female:provided_support_actor_cwp 
    #+ is_female:provided_support_partner_cwp 
  | coupleID, 
  
  weights = varIdent(form = ~ 1 | gender), # heterogeneous residual variances
  corr = corCompSymm(form = ~ 1 | coupleID/diaryday), # compound symmetry (partner's values on each day)
  
  data = df_long_apim, 
  control = list(maxIter=1000)
)


summary(model1)

formula <- bf(
  closeness ~ 0 + is_male + is_female +
    
    # Time Slopes
    is_male:diaryday_c +
    is_female:diaryday_c +
    
    # Between-Person APIM
    is_male:provided_support_actor_cbp + 
    is_male:provided_support_partner_cbp +
    
    is_female:provided_support_actor_cbp + 
    is_female:provided_support_partner_cbp +
  
    # Within-Person APIM
    is_male:provided_support_actor_cwp + 
    is_male:provided_support_partner_cwp +
    
    is_female:provided_support_actor_cwp + 
    is_female:provided_support_partner_cwp +
  
  # Accounting for non-independence between partners' means and trajectories 
  # and effect sensitivities via random effects:
  (0 + is_male + is_female
    + is_male:diaryday_c + is_female:diaryday_c
  | coupleID) +
    
    
  # Accounting for daily non-independence (mimicking nlme's corCompSym)
  # modelled via common 'shocks' with a random intercept:
  (1 | coupleID:diaryday)

  # heteroscedastic residuals
  , sigma ~ 0 + is_male + is_female 
)

priors <- c(
  prior(normal(4, 2), class = "b", coef = "is_male"), # male intercept
  prior(normal(4, 2), class = "b", coef = "is_female"), # female intercept
  prior(normal(0, 2), class = "b") # other beta coefficients
  # Other priors can be specified, but brms choosed defaults well. 
  # Defaults can be seen via default_prior(formula, data)
)

model_dist_apim_long_simple <- brm(
  formula = formula, 
  data = df_long_apim,
  family = gaussian(link = identity), # student() is also often a nice robust option. 
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_dist_apim_long_simpel') # Cache the model
)

summary(model_dist_apim_long_simple)

    ... +
    ar(time = diaryday, gr = coupleID:userID, p = 1),
    
  sigma ~ 0 + is_male + is_female
)

loo1 <- loo(model_dist_apim_long_simple)
loo2 <- loo(model_dist_apim_long_ar)

loo::pareto_k_table(loo1)

loo::pareto_k_table(loo2)

a <- loo_compare(loo1, loo2)
print(a)

report::report(a)

formula <- bf(
  closeness ~ 
    
    # Intercepts
    0 + is_male + is_female +
    
    # Time Slopes
    is_male:diaryday_c +
    is_female:diaryday_c +
    
    # Between-Person APIM
    is_male:provided_support_actor_cbp + 
    is_male:provided_support_partner_cbp +
    
    is_female:provided_support_actor_cbp + 
    is_female:provided_support_partner_cbp +
  
    # Within-Person APIM
    is_male:provided_support_actor_cwp + 
    is_male:provided_support_partner_cwp +
    
    is_female:provided_support_actor_cwp + 
    is_female:provided_support_partner_cwp +
    
    # Accounting for non-independence between partners' means and trajectories 
    # and effect sensitivities via random effects:
    (0 + is_male + is_female +  
       
       is_male:diaryday_c + is_female:diaryday_c +
       
       is_male:provided_support_actor_cwp + 
       is_male:provided_support_partner_cwp +
       is_female:provided_support_actor_cwp + 
       is_female:provided_support_partner_cwp 
    | coupleID ) + 
    
    
    # Accounting for daily non-independence
    (1 | coupleID:diaryday), 
  
  # No more AR1. You could also test MA or ARMA. 
    
  sigma ~ 0 + is_male + is_female   # heteroscedastic residuals
)

priors <- c(
  prior(normal(4, 2), class = "b", coef = "is_male"),
  prior(normal(4, 2), class = "b", coef = "is_female"),
  prior(normal(0, 2), class = "b")
)


model_dist_apim_long_complex <- brm(
  formula = formula, 
  data = df_long_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_dist_apim_long_complex') # Cache the model
)

loo1 <- loo(model_dist_apim_long_simple)
loo2 <- loo(model_dist_apim_long_complex)

loo::pareto_k_table(loo1)

loo::pareto_k_table(loo2)

a <- loo_compare(loo1, loo2)
print(a)

report::report(a)

print_df(head(df_apim))

formula <- bf(
  satisfaction ~ 1 + 
    communication_actor_gmc + communication_partner_gmc + 
    
    # Option 1: A single couple level random intercept.
    # In the cross sectional case this is the maximal random effects structure and
    # we cannot distinguish each partners variance from noise. 
    # (1 | coupleID)
    
    # Option 2: using a compound symmetry residual structure. 
    cosy(gr = coupleID)
  
  # Note: no need to model separate sigmas for each partner.
  # Homogeneous residual variance is estimated. 
  # Implied: sigma = ~ 1
)

priors <- c(
  prior(normal(2, 10), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "sigma"),
  prior(beta(1, 3), class = "cosy")
)

model_ind_apim <- brm(
  formula = formula, 
  data = df_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_ind_apim') # Cache the model
)

a <- loo_compare(
  loo(model_ind_apim), 
  loo(model_dist_apim_b)
)
print(a)

report::report(a)

df_dim <- df_apim2 %>%
  group_by(coupleID) %>%
  mutate(
    communication_cm = mean(communication, na.rm = TRUE),
    communication_cwc= communication - communication_cm
  ) %>%
  ungroup() %>%
  mutate(
    communication_cbc= communication_cm - mean(communication_cm, na.rm = TRUE)
  ) %>% 
  dplyr::select(c('userID', 'coupleID', 'satisfaction', 'communication_cwc', 'communication_cbc'))

print_df(head(df_dim))

formula <- bf(
  satisfaction ~ 1 + 
    communication_cbc+ communication_cwc + 
    cosy(gr = coupleID) 
)

priors <- c(
  prior(normal(2, 10), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "sigma"),
  prior(beta(1, 3), class = "cosy")
)

model_ind_dim <- brm(
  formula = formula, 
  data = df_dim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_ind_dim') # Cache the model
)

a <- hypothesis(
  model_ind_apim, 
  "communication_actor_gmc + communication_partner_gmc = 0"
)

a$hypothesis[,c(2,3,4,5)]

df_long_dim <- df_long_apim2 %>%
  mutate(
    
    # Between Person Level DIM (centering between couples and within couple between person)
    provided_support_cbc = (provided_support_actor_cbp + provided_support_partner_cbp) / 2,
    provided_support_cwcbp = (provided_support_actor_cbp - provided_support_partner_cbp) / 2,
    
    # Within Person Level DIM (couple mean and deviation on each day)
    provided_support_cwp_mean = (provided_support_actor_cwp + provided_support_partner_cwp) / 2,
    provided_support_cwp_halfdiff = (provided_support_actor_cwp - provided_support_partner_cwp) / 2
  ) 

df_long_dim <- df_long_dim %>%
  group_by(coupleID) %>%
  mutate(
    base = ifelse(userID == min(userID), 1, -1),
    flip = 1 - 2 * rbinom(1, 1, 0.5),   # yields +1 or -1 once per couple
    Idiff = base * flip
  ) %>%
  ungroup() %>%
  relocate(Idiff, .after = coupleID) %>%
  dplyr::select(-base, -flip)


print_df(print_couple_preview(df_long_dim, couple_id = "31"))  

formula <- bf(
  closeness ~ 1 + 
    
    diaryday_c +
  
    # Within-person APIM
    provided_support_actor_cwp + 
    provided_support_partner_cwp +
    # Equivalent to within-person DIM:
    # provided_support_cwp_mean + 
    # provided_support_cwp_halfdiff +
  
    # Between-person APIM
    provided_support_actor_cbp +
    provided_support_partner_cbp +
    # Equivalent to between-person DIM
    # provided_support_cwcbp+
    # provided_support_cbc +
    
    # Dyad-Level intercept and slopes for time-varying predictors
    (1 + diaryday_c + provided_support_actor_cwp + provided_support_partner_cwp | coupleID) +
    # Both partners' deviations from these dyad-level means and slopes.
    # Note: separate random effects block to make them uncorrelated form dyad-level REs. 
    (0 + Idiff + 
         I(Idiff * diaryday_c) + 
         I(Idiff * provided_support_actor_cwp) + 
         I(Idiff * provided_support_partner_cwp) | coupleID) +
    
    # Accounting for same-day shocks/coupling
    (1 | coupleID:diaryday)
    
    # Autocorrelated residuals
    # Due to bad fit in the distinguishable case, we omit AR1 from the start.
    # ar(time = diaryday, gr = coupleID:userID, p = 1)
  
  # Again, no need to model heterogeneous residual variances (sigma)
  # Implied: sigma = ~ 1
)

priors <- c(
  prior(normal(4, 2), class = "Intercept"),
  prior(normal(0, 2), class = "b")
  
  # We could set priors for other things, but brms sets good default priors
  #prior(student_t(3, 0, 1.5), class = "sd"),
  #prior(lkj(2), class = "cor"), # correlation prior for random effect matrix
  #prior(student_t(3, 0, 1.5), class = "sigma")
)

model_apim_ind_long <- brm(
  formula = formula, 
  data = df_long_dim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'model_apim_ind_long_apim')
)

# Custom function 
rotate_apim_covariance <- function(
    fit,
    gr = "coupleID",
    Idiff = "Idiff"
) {
  random_effects <- fit$ranef
  random_effects <- random_effects[random_effects$group == gr,]

  SUM <- random_effects$coef[!grepl(Idiff, random_effects$coef)]
  DIFF <- random_effects$coef[grepl(Idiff, random_effects$coef)]
  
  varcor <- VarCorr(fit)[[gr]]
  
  within_person_covariance_matrix <- varcor$cov[SUM,'Estimate',SUM] + varcor$cov[DIFF,'Estimate',DIFF]
  cross_person_covariance_matrix <- varcor$cov[SUM,'Estimate',SUM] - varcor$cov[DIFF,'Estimate',DIFF]
  
  p <- nrow(within_person_covariance_matrix)
  
  Full <- rbind(
    cbind(within_person_covariance_matrix,  cross_person_covariance_matrix),
    cbind(cross_person_covariance_matrix, within_person_covariance_matrix)
  )
  
  # nice labels
  base <- SUM
  rn <- c(paste0("PartnerA_", base), paste0("PartnerB_", base))
  colnames(Full) <- rn
  rownames(Full) <- rn
  
  return(list(within_person_covariance_matrix=within_person_covariance_matrix, cross_person_covariance_matrix=cross_person_covariance_matrix, full_covariance_matrix=Full))
}

a <- rotate_apim_covariance(model_apim_ind_long)

print_df(a$within_person_covariance_matrix)

print_df(a$cross_person_covariance_matrix)

print_df(a$full_covariance_matrix)

# Custom function
sd_cor <- function(Sigma) {
  sds  <- sqrt(diag(Sigma))
  cors <- cov2cor(Sigma)  
  list(sd = sds, cor = cors)
}

full <- sd_cor(a$full_covariance_matrix)

print_df(as.data.frame(round(full$sd, 3)))

print_df(as.data.frame(round(full$cor, 3)))